### Regret Aversion (Bell, 1982; Loomes and Sugden, 1982) function library
## Appendix Theory: A.4.5

MaxLikeSR_regretAversion_LOO_crossVal = function(data, sub_sample, init, df_regret_probs){
  # Leave-one-out Cross Validation (LOO CV)
  start <- Sys.time()  
  data = sample_SR(data,sub_sample)
  index = data$studentID                                                            
  data = split(data,f = index)  
  
  dfs_LOOCV = lapply(data, LOOCV_train_test_split)
  results = lapply(dfs_LOOCV, LOO_crossVal_train_predict_regretAver, init, df_regret_probs)
  alpha_coef = unlist(lapply(results, {function(x) mean(x[["alpha.coef"]])}), use.names = FALSE)
  alpha_coef_std = unlist(lapply(results, {function(x) sqrt(var(x[["alpha.coef"]]))}), use.names = FALSE)
  beta_coef = unlist(lapply(results, {function(x) mean(x[["beta.coef"]])}), use.names = FALSE)
  beta_coef_std = unlist(lapply(results, {function(x) sqrt(var(x[["beta.coef"]]))}), use.names = FALSE)
  sigma_coef = unlist(lapply(results, {function(x) mean(x[["sigma.coef"]])}), use.names = FALSE)
  sigma_coef_std = unlist(lapply(results, {function(x) sqrt(var(x[["sigma.coef"]]))}), use.names = FALSE)
  
  score = unlist(lapply(results, {function(x) sum(x$accuracy)}), use.names = FALSE)
  results_scores = as.data.frame(cbind(studentID = unique(index), score = score, alpha_mean = alpha_coef, beta_mean = beta_coef, sigma_mean = sigma_coef,
                                       alpha_std = alpha_coef_std, beta_std = beta_coef_std, sigma_std = sigma_coef_std))
  
  print("Regret Aversion - Leave-One-Out Test Accuracy Summary:")
  print(paste("Average:", mean(results_scores$score)))
  print(paste("Standard error:", sqrt(var(results_scores$score)/length(results_scores$score))))
  
  end = Sys.time()                                                                  
  timer = difftime(end,start,units = "secs")
  print(paste("Time taken for MLE algorithim and Cross-Validation:",seconds_to_period(round(timer[[1]],2))))
  return(list(estimations = results, summary = results_scores))
}



LOOCV_train_test_split = function(data){
  result = list()
  for(i in 1:length(data$qnum)){
    test_df = data[i, ]
    train_df = data[-i, ]
    result[[i]] <- list(train = train_df, test = test_df)
  }
  result
}



LOO_crossVal_train_predict_regretAver = function(LOO_train_test_folds, init, df_regret_probs){
  results = bind_rows(lapply(LOO_train_test_folds, LOO_train_test_regret, init, df_regret_probs))
  results
}



LOO_train_test_regret = function(train_test_dfs, init, df_regret_probs){
  train_df = train_test_dfs[["train"]]
  test_df = train_test_dfs[["test"]]
  
  results = wrapMLE_regret(train_df, init, log_lik_regretAver, df_regret_probs)
  results_summary = stat_sum_regret(results)
  results_summary$studentID = test_df$studentID
  
  param = list(alpha = results_summary$alpha.coef, beta = results_summary$beta.coef, sigma = results_summary$sigma.coef)
  test_predict_results = LOO_test_predict_regret(test_df, param, calc_util_regretAver, df_regret_probs)
  results_summary = cbind(results_summary, test_qnum = test_df$qnum, test_choice = test_df$choice, test_predict_results)
  results_summary
}





LOO_test_predict_regret = function(test_df, param, util_fun, df_regret_probs){
  lottery_data = list(l1 = list(prob = test_df[["prob.1"]], lower = test_df[["lower.1"]], upper = test_df[["upper.1"]]),
                      l2 = list(prob = test_df[["prob.2"]], lower = test_df[["lower.2"]], upper = test_df[["upper.2"]]),
                      l3 = list(prob = test_df[["prob.3"]], lower = test_df[["lower.3"]], upper = test_df[["upper.3"]]),
                      l4 = list(prob = test_df[["prob.4"]], lower = test_df[["lower.4"]], upper = test_df[["upper.4"]]),
                      l5 = list(prob = test_df[["prob.5"]], lower = test_df[["lower.5"]], upper = test_df[["upper.5"]]))
  choice = test_df[["choice"]]
  qnum = test_df[["qnum"]]
  lottery_utility = unlist(lapply(lottery_data, util_fun, param, lottery_data, choice, qnum, df_regret_probs))
  predicted_choice = unname(which.max(lottery_utility))
  accuracy = if_else(predicted_choice == test_df$choice, 1, 0)
  list(pred_choice = predicted_choice, accuracy = accuracy)
}



calc_util_regretAver = function(df, param, lottery_data, choice, qnum, df_regret_probs){
  df[["prob"]]*((55+df[["upper"]])^param[["sigma"]]) + (1-df[["prob"]])*((55+df[["lower"]])^param[["sigma"]]) +
  df[["prob"]]*k_rule_fun((55+df[["upper"]])^param[["sigma"]] - calc_max_regret_num(choice, qnum, 'h', df_regret_probs, lottery_data, param[["sigma"]]), param[["alpha"]], param[["beta"]]) + 
  (1-df[["prob"]])*k_rule_fun((55+df[["lower"]])^param[["sigma"]] - calc_max_regret_num(choice, qnum, 'l', df_regret_probs, lottery_data, param[["sigma"]]), param[["alpha"]], param[["beta"]])
}


k_rule_fun = function(rule, alpha, beta){
  if(rule>=0){alpha*(rule^beta)}
  else{-alpha*((-rule)^beta)}
}


calc_max_regret_num = function(choice, qnum, h_or_l, df_regret_probs, lottery_data, sigma){
  range = 1:5
  range = range[-choice]
  result = list()
  for(i in range){
    upper = lottery_data[[i]][["upper"]]
    lower = lottery_data[[i]][["lower"]]
    q = get_q(choice, qnum, h_or_l, i, df_regret_probs)
    res = (q*((55+upper)^sigma)) + ((1-q)*((55+lower)^sigma))
    result[[length(result)+1]] = res
  }
  max_result = max(unlist(result))
  max_result
}



get_q = function(choice, qnum, h_or_l, j, df_regret_probs){
  if(h_or_l=='h'){start=2}
  else if(h_or_l=='l'){start=7}
  q = df_regret_probs[(df_regret_probs$qnum==qnum & df_regret_probs$choiceJ==j), start+choice]
  as.numeric(q)
}




wrapMLE_regret = function(data,init,logLikFun, df_regret_probs){
  regretAver_MaxLike_res = mle2(minuslogl = logLikFun, data = list(data=data, df_regret = df_regret_probs), start = init,
                               lower = c(alpha=0, beta=1, sigma=0.01), upper = c(alpha=2, beta=2, sigma=1))
  regretAver_MaxLike_res
}





stat_sum_regret = function(stat){
  results = data.frame("studentID"=NA,"alpha.coef"=NA, "beta.coef"=NA, "sigma.coef"=NA,
                       "std.alpha"=NA, "std.beta"=NA, "std.sigma"=NA,
                       "p.val.alpha"=NA, "p.val.beta"=NA, "p.val.sigma"=NA,
                       "logLikelihood"=NA,"AIC"=NA)
  summary = summary(stat)
  results["studentID"]  = NA
  results["alpha.coef"]  = summary@coef[1,1]
  results["beta.coef"]   = summary@coef[2,1]
  results["sigma.coef"]  = summary@coef[3,1]
  results["std.alpha"]   = summary@coef[1,2]
  results["std.beta"]    = summary@coef[2,2]
  results["std.sigma"]   = summary@coef[3,2]
  results["p.val.alpha"] = summary@coef[1,4]
  results["p.val.beta"]  = summary@coef[2,4]
  results["p.val.sigma"] = summary@coef[3,4]
  results["logLikelihood"] = logLik(stat)
  results["AIC"]  = AIC(stat)
  return(results)
}





log_lik_regretAver <- function(alpha, beta, sigma, data, df_regret){
  #calculating denominator of likelihood of a choice
    denomCalc = function(data, df_regret){
      lottery_data = create_lottery_data(data)
      
      exp(utility_fun_regret(lottery_data[[1]], lottery_data, choice = 1, data[["qnum"]], df_regret, alpha, beta, sigma)) +
      exp(utility_fun_regret(lottery_data[[2]], lottery_data, choice = 2, data[["qnum"]], df_regret, alpha, beta, sigma)) +
      exp(utility_fun_regret(lottery_data[[3]], lottery_data, choice = 3, data[["qnum"]], df_regret, alpha, beta, sigma)) +
      exp(utility_fun_regret(lottery_data[[4]], lottery_data, choice = 4, data[["qnum"]], df_regret, alpha, beta, sigma)) +
      exp(utility_fun_regret(lottery_data[[5]], lottery_data, choice = 5, data[["qnum"]], df_regret, alpha, beta, sigma))
    }
    
    
    denomRes = apply(data, 1, denomCalc, df_regret) #vector of denominator components for each student
    
    # calculate numerator of likelihood of a choice
    numerCalc = function(data, df_regret){
      lottery_data = create_lottery_data(data)
      
      j = data[["choice"]]
      numercalc = exp(data[[13+j]]*((55+data[[8+j]])^sigma) + (1-data[[13+j]])*((55+data[[3+j]])^sigma) +
                      data[[13+j]]*k_rule_fun((55+data[[8+j]])^sigma - calc_max_regret_num(data[["choice"]], data[["qnum"]], 'h', df_regret, lottery_data, sigma), alpha, beta) +
                      (1-data[[13+j]])*k_rule_fun((55+data[[3+j]])^sigma - calc_max_regret_num(data[["choice"]], data[["qnum"]], 'l', df_regret, lottery_data, sigma), alpha, beta))
      if(is.infinite(numercalc)){browser()}
      else if(is.na(numercalc)){browser()}
      else if(numercalc==0){browser()}
      numercalc
    }
    
    numerRes = apply(data, 1, numerCalc, df_regret)
    
    #calculating log likelihood
    log_lik_i <-  log(numerRes/denomRes)
    
    -sum(unlist(log_lik_i))
}




utility_fun_regret = function(data_i, data_all, choice, qnum, df_regret_probs, alpha, beta, sigma){
  prob = data_i[["prob"]]
  upper = data_i[["upper"]]
  lower = data_i[["lower"]]
  
  res = prob*((55+upper)^sigma) + (1-prob)*((55+lower)^sigma) +
  prob*k_rule_fun((55+upper)^sigma - calc_max_regret_num(choice, qnum, 'h', df_regret_probs, data_all, sigma), alpha, beta) + 
  (1-prob)*k_rule_fun((55+lower)^sigma - calc_max_regret_num(choice, qnum, 'l', df_regret_probs, data_all, sigma), alpha, beta)
  res
}



create_lottery_data = function(data){
  lottery_data = list(l1 = list(prob = data[["prob.1"]], lower = data[["lower.1"]], upper = data[["upper.1"]]),
                      l2 = list(prob = data[["prob.2"]], lower = data[["lower.2"]], upper = data[["upper.2"]]),
                      l3 = list(prob = data[["prob.3"]], lower = data[["lower.3"]], upper = data[["upper.3"]]),
                      l4 = list(prob = data[["prob.4"]], lower = data[["lower.4"]], upper = data[["upper.4"]]),
                      l5 = list(prob = data[["prob.5"]], lower = data[["lower.5"]], upper = data[["upper.5"]]))

  lottery_data  
}


sample_SR = function(data,sub_sample){
  if(sub_sample == "full"){data = data}
  if(sub_sample == "part1"){data = data %>% filter(qnum<19)}
  if(sub_sample == "part2"){data = data %>%  filter(qnum>18)}
  return(data)
}

